package top.atcra.rpc.protection.loadBalancer.selector.impl;

import lombok.extern.slf4j.Slf4j;
import top.atcra.rpc.protection.loadBalancer.selector.ServiceSelector;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

@Slf4j
public class ConsistentHashSelector extends ServiceSelector {

    protected static final int PER_NODE_VIRTUAL_NODE_COUNT = 64;
    // hash环，将虚拟结点映射到真实结点
    protected final SortedMap<Integer, InetSocketAddress> ring = new TreeMap<>();

    public ConsistentHashSelector(List<InetSocketAddress> addresses) {
        super(addresses);
        for (InetSocketAddress address : addresses) {
            insertHostIntoHashCircle(address);
        }
    }

    /**
     * 将每一个服务器节点地址挂载到hash环上
     *
     * @param inetSocketAddress 服务器节点地址
     */
    private void insertHostIntoHashCircle(InetSocketAddress inetSocketAddress) {
        // 为每一个节点生成虚拟节点进行挂载
        for (int i = 0; i < PER_NODE_VIRTUAL_NODE_COUNT; i++) {
            int hash = (inetSocketAddress.toString() + "#" + i).hashCode();
            // 挂载到hash环上
            ring.put(hash, inetSocketAddress);

            if (log.isDebugEnabled()) {
                log.debug("服务器地址为【{}】,hash为【{}】的虚拟节点已经挂载到了hash环上.", inetSocketAddress, hash);
            }
        }
    }

    @Override
    public InetSocketAddress getAddress(Object key) {
        if (ring.isEmpty()) return null;
        int hash = key.hashCode();
        // 找到第一个大于等于hash的子树
        var rtree = ring.tailMap(hash);
        if (rtree.isEmpty()) {
            return ring.firstEntry().getValue();
        } else {
            return rtree.firstEntry().getValue();
        }
    }
}
